import numpy as np
import matplotlib.pyplot as plt

class DoubleMoonGenerator:
    def __init__(self, N, r, w, d):
        self.N = N  # 数据点数量
        self.r = r  # 月亮半径
        self.w = w  # 月亮宽度
        self.d = d  # 两个月亮之间的距离
        self.data = None  # 用于存储生成的数据

    def generate_moon(self):
        data = np.empty([2 * self.N, 3])  # 添加一列用于标签
        count = 0

        while count < self.N:
            x = 2 * (self.r + self.w / 2) * (np.random.random() - 0.5)
            y = (self.r + self.w / 2) * np.random.random()
            distance = np.sqrt(x * x + y * y)

            if (distance < self.r + self.w / 2) and (distance > self.r - self.w / 2):
                data[count] = [x, y, 1]  # 标记为上半月，标签为 1
                count += 1

        for i in range(self.N):
            data[self.N + i] = [data[i, 0] + self.r, -data[i, 1] - self.d, -1]  # 标记为下半月，标签为 -1

        self.data = data
        return data

    def save_to_file(self, filename):
        if self.data is not None:
            np.savetxt(filename, self.data, delimiter=',')  # 保存数据到文本文件

class LeastSquaresEstimation:
    def __init__(self, data):
        self.data = data

    def fit(self):
        X = self.data[:, :2]  # 特征
        y = self.data[:, 2]   # 标签

        # 添加偏置项
        X_b = np.c_[np.ones((len(X), 1)), X]

        # 最小二乘估计
        self.theta = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y)

    def predict(self, X):
        X_b = np.c_[np.ones((len(X), 1)), X]
        return X_b.dot(self.theta)

class MaximumLikelihoodEstimation:
    def __init__(self, data):
        self.data = data

    def fit(self):
        X = self.data[:, :2]  # 特征
        y = self.data[:, 2]   # 标签

        # 添加偏置项
        X_b = np.c_[np.ones((len(X), 1)), X]

        # 最大似然估计
        self.theta = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y)

    def predict(self, X):
        X_b = np.c_[np.ones((len(X), 1)), X]
        return X_b.dot(self.theta)

class MaximumAPosterioriEstimation:
    def __init__(self, data, alpha=0.1):
        self.data = data
        self.alpha = alpha

    def fit(self):
        X = self.data[:, :2]  # 特征
        y = self.data[:, 2]   # 标签

        # 添加偏置项
        X_b = np.c_[np.ones((len(X), 1)), X]

        # 最大后验估计
        eye = np.eye(X_b.shape[1])
        eye[0, 0] = 0  # 不对截距项进行正则化
        self.theta = np.linalg.inv(X_b.T.dot(X_b) + self.alpha * eye).dot(X_b.T).dot(y)

    def predict(self, X):
        X_b = np.c_[np.ones((len(X), 1)), X]
        return X_b.dot(self.theta)

if __name__ == "__main__":
    N = 1000
    r = 10
    w = 6
    d = 1
    moon_generator = DoubleMoonGenerator(N, r, w, d)
    data = moon_generator.generate_moon()
    moon1 = data[:N]
    moon2 = data[N:]
    plt.figure(figsize=(8, 6))
    plt.scatter(moon1[:, 0], moon1[:, 1], color='blue', label='Moon 1 (Upper)')
    plt.scatter(moon2[:, 0], moon2[:, 1], color='red', label='Moon 2 (Lower)')
    # 将数据保存到带有标签的文本文件中
    moon_generator.save_to_file('double_moon_data_with_labels.txt')

    method = "LSE"  # 设置要执行的方法:LSE|MLE|MAP
    '''设置要执行的方法
    @param  LSE:    最小二乘估计
    @param  MLE:    最大似然估计
    @param  MAP:    最大后验估计
    '''

    if method == "LSE": 
        # 创建 LeastSquaresEstimation 实例
        ls_estimator = LeastSquaresEstimation(data)

        # 拟合数据
        ls_estimator.fit()

        # 输出模型参数
        print(f"模型参数：{ls_estimator.theta}")

        # 预测双月数据的分类
        moon1_predicted = ls_estimator.predict(data[:N, :2])
        moon2_predicted = ls_estimator.predict(data[N:, :2])

        # 展示双月数据和决策边界
        plt.figure(figsize=(8, 6))
        plt.scatter(moon1[:, 0], moon1[:, 1], color='blue', label='Moon 1 (Upper)')
        plt.scatter(moon2[:, 0], moon2[:, 1], color='red', label='Moon 2 (Lower)')

        # 画出决策直线
        x_values = np.linspace(np.min(data[:, 0]), np.max(data[:, 0]), 100)
        y_values3 = ls_estimator.theta[0] + ls_estimator.theta[1] * x_values
        plt.plot(x_values, y_values3, color='green', label='Decision Boundary (Least Squares)')
        plt.title('LSE Classification on Moon Dataset') 
        # 保存最小二乘估计的决策边界图
        plt.savefig('decision_boundary_LSE.png')
    
    elif method == "MLE":
        # 创建 MaximumLikelihoodEstimation 实例
        mle = MaximumLikelihoodEstimation(data)

        # 拟合数据
        mle.fit()

        # 输出模型参数
        print(f"模型参数：{mle.theta}")

        # 预测双月数据的分类
        moon1_predicted = mle.predict(data[:N, :2])
        moon2_predicted = mle.predict(data[N:, :2])

        # 展示双月数据和决策边界
        plt.figure(figsize=(8, 6))
        plt.scatter(moon1[:, 0], moon1[:, 1], color='blue', label='Moon 1 (Upper)')
        plt.scatter(moon2[:, 0], moon2[:, 1], color='red', label='Moon 2 (Lower)')

        # 画出决策直线
        x_values = np.linspace(np.min(data[:, 0]), np.max(data[:, 0]), 100)
        y_values1 = mle.theta[0] + mle.theta[1] * x_values
        plt.plot(x_values, y_values1, color='green', label='Decision Boundary (MLE)')
        plt.title('MLE Classification on Moon Dataset')
        # 保存最大似然估计的决策边界图
        plt.savefig('decision_boundary_MLE.png')

    elif method == "MAP":
        # 创建 MaximumAPosterioriEstimation 实例
        map_estimator = MaximumAPosterioriEstimation(data)

        # 拟合数据
        map_estimator.fit()

        # 输出模型参数
        print(f"模型参数：{map_estimator.theta}")

        # 预测双月数据的分类
        moon1_predicted = map_estimator.predict(data[:N, :2])
        moon2_predicted = map_estimator.predict(data[N:, :2])

        # 展示双月数据和决策边界
        plt.figure(figsize=(8, 6))
        plt.scatter(moon1[:, 0], moon1[:, 1], color='blue', label='Moon 1 (Upper)')
        plt.scatter(moon2[:, 0], moon2[:, 1], color='red', label='Moon 2 (Lower)')

        # 画出决策直线
        x_values = np.linspace(np.min(data[:, 0]), np.max(data[:, 0]), 100)
        y_values2 = map_estimator.theta[0] + map_estimator.theta[1] * x_values
        plt.plot(x_values, y_values2, color='green', label='Decision Boundary (MAP)')
        plt.title('MAP Classification on Moon Dataset')
        # 保存最大后验估计的决策边界图
        plt.savefig('decision_boundary_MAP.png')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend()
    plt.show()